from fastapi import FastAPI, HTTPException, status, BackgroundTasks, Depends
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import Optional, Dict, List, Any
import uuid
import time
import asyncio
from enum import Enum
import os
from dotenv import load_dotenv
import logging

from autodl_client import WanInstanceManager, AutoDLClient

# Load environment variables
load_dotenv()

app = FastAPI(title="Wan 2.2 Video Generation API",
              description="API for generating videos using Wan 2.2 with AutoDL elastic deployment",
              version="1.0.0")

# CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Enums
class TaskStatus(str, Enum):
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"

class GenerationMode(str, Enum):
    T2V = "t2v"  # Text to Video
    I2V = "i2v"  # Image to Video
    TI2V = "ti2v"  # Text and Image to Video

# Models
class TaskCreate(BaseModel):
    mode: GenerationMode
    params: Dict[str, Any] = Field(
        ...,
        example={
            "prompt": "A beautiful sunset over the mountains",
            "duration": 5,
            "resolution": "720p",
            "fps": 24,
            "style": "cinematic"
        }
    )

class TaskResponse(BaseModel):
    task_id: str
    status: TaskStatus
    result_url: Optional[str] = None
    error: Optional[str] = None

# Initialize AutoDL client
autodl_client = AutoDLClient(
    api_key=os.getenv("AUTODL_API_KEY"),
    endpoint=os.getenv("AUTODL_ENDPOINT", "https://api.autodl.com")
)

# Initialize Wan instance manager
wan_manager = WanInstanceManager(
    autodl_client=autodl_client,
    max_instances=int(os.getenv("MAX_NODES", 2))
)

# In-memory storage (replace with database in production)
tasks_db = {}
task_queue = []
active_tasks = 0

# AutoDL configuration
AUTODL_CONFIG = {
    "api_key": os.getenv("AUTODL_API_KEY"),
    "endpoint": os.getenv("AUTODL_ENDPOINT", "https://api.autodl.com"),
    "wan_image": "wan-2.2"
}

# API Endpoints
@app.post("/api/tasks/", response_model=TaskResponse, status_code=status.HTTP_201_CREATED)
async def create_task(background_tasks: BackgroundTasks, task: TaskCreate):
    task_id = str(uuid.uuid4())
    
    # Create task entry
    task_data = {
        "task_id": task_id,
        "status": TaskStatus.PENDING,
        "mode": task.mode,
        "params": task.params,
        "created_at": time.time(),
        "started_at": None,
        "completed_at": None,
        "result_url": None,
        "error": None,
        "instance_id": None,
        "progress": 0
    }
    
    # Add to database
    tasks_db[task_id] = task_data
    
    # Start processing in background
    background_tasks.add_task(process_task_queue)
    
    return task_data

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    try:
        # Check Redis connection if configured
        if 'redis' in globals():
            await redis.ping()
            
        # Check database connection if configured
        if 'db' in globals():
            await db.execute("SELECT 1")
            
        return {
            "status": "healthy",
            "version": "1.0.0",
            "services": {
                "database": "connected",
                "redis": "connected" if 'redis' in globals() else "not configured",
                "autodl": "connected" if autodl_client else "not configured"
            }
        }
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail=f"Service unavailable: {str(e)}"
        )

@app.get("/api/tasks/{task_id}", response_model=TaskResponse)
async def get_task_status(task_id: str):
    task = tasks_db.get(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    
    # If task is processing, check instance status
    if task["status"] == TaskStatus.PROCESSING and task.get("instance_id"):
        try:
            instance = autodl_client.get_instance(task["instance_id"])
            if instance.get("status") == "error":
                task["status"] = TaskStatus.FAILED
                task["error"] = "Instance error"
                task["completed_at"] = time.time()
        except Exception as e:
            logger.warning(f"Error checking instance status: {e}")
    
    return task

# Task processing
async def process_task_queue():
    global active_tasks
    
    # Process tasks until queue is empty or max instances reached
    while task_queue:
        # Check if we can start more instances
        if not wan_manager.can_scale_up():
            break
            
        # Get the next task from the queue
        task_id = task_queue.pop(0)
        task = tasks_db.get(task_id)
        
        if task and task["status"] == TaskStatus.PENDING:
            # Try to start a new instance
            instance = wan_manager.start_instance()
            if not instance:
                # If we can't start a new instance, put task back in queue
                task_queue.insert(0, task_id)
                break
                
            # Update task status
            task["status"] = TaskStatus.PROCESSING
            task["started_at"] = time.time()
            task["instance_id"] = instance["id"]
            
            # Start processing task in background
            active_tasks += 1
            asyncio.create_task(process_task(task))

async def process_task(task):
    """Process a video generation task"""
    global active_tasks
    
    try:
        logger.info(f"Starting task {task['task_id']} in {task['mode']} mode")
        
        # Get the instance for this task
        instance_id = task.get("instance_id")
        if not instance_id:
            raise ValueError("No instance assigned to task")
            
        # Get instance details
        instance = autodl_client.get_instance(instance_id)
        if not instance or instance.get("status") != "running":
            raise ValueError("Instance not available")
            
        # Get instance endpoint
        instance_ip = instance.get("ip_address")
        if not instance_ip:
            raise ValueError("Instance has no IP address")
            
        # Prepare request to Wan instance
        wan_endpoint = f"http://{instance_ip}:8000"
        
        # Call the appropriate endpoint based on mode
        if task["mode"] == GenerationMode.T2V:
            await call_wan_t2v(wan_endpoint, task)
        elif task["mode"] == GenerationMode.I2V:
            await call_wan_i2v(wan_endpoint, task)
        elif task["mode"] == GenerationMode.TI2V:
            await call_wan_ti2v(wan_endpoint, task)
        else:
            raise ValueError(f"Unsupported mode: {task['mode']}")
        
        # Mark task as completed
        task["status"] = TaskStatus.COMPLETED
        task["completed_at"] = time.time()
        task["result_url"] = f"{os.getenv('STORAGE_BASE_URL')}/{task['task_id']}.mp4"
        
        logger.info(f"Completed task {task['task_id']}")
        
    except Exception as e:
        logger.error(f"Error processing task {task.get('task_id')}: {str(e)}")
        task["status"] = TaskStatus.FAILED
        task["error"] = str(e)
        task["completed_at"] = time.time()
        
    finally:
        # Clean up
        active_tasks -= 1
        
        # Stop the instance if no more tasks
        if instance_id and active_tasks == 0:
            logger.info(f"No active tasks, stopping instance {instance_id}")
            autodl_client.stop_instance(instance_id)
        
        # Process next task in queue
        asyncio.create_task(process_task_queue())

async def call_wan_t2v(wan_endpoint: str, task: Dict):
    """Call Wan T2V endpoint"""
    url = f"{wan_endpoint}/api/t2v"
    params = task["params"]
    
    # Add required parameters
    payload = {
        "prompt": params.get("prompt", ""),
        "duration": params.get("duration", 5),
        "resolution": params.get("resolution", "720p"),
        "fps": params.get("fps", 24),
        "style": params.get("style", "realistic"),
        "task_id": task["task_id"]
    }
    
    # Make the request
    async with aiohttp.ClientSession() as session:
        async with session.post(url, json=payload) as response:
            if response.status != 200:
                error = await response.text()
                raise Exception(f"Wan API error: {error}")
            
            # Stream the response to update progress
            async for chunk in response.content.iter_any():
                try:
                    progress_data = json.loads(chunk)
                    if "progress" in progress_data:
                        task["progress"] = progress_data["progress"]
                except:
                    continue

async def call_wan_i2v(wan_endpoint: str, task: Dict):
    """Call Wan I2V endpoint"""
    url = f"{wan_endpoint}/api/i2v"
    params = task["params"]
    
    # Add required parameters
    payload = {
        "image_url": params.get("image_url"),
        "duration": params.get("duration", 5),
        "fps": params.get("fps", 24),
        "task_id": task["task_id"]
    }
    
    # Make the request
    async with aiohttp.ClientSession() as session:
        async with session.post(url, json=payload) as response:
            if response.status != 200:
                error = await response.text()
                raise Exception(f"Wan API error: {error}")
            
            # Stream the response to update progress
            async for chunk in response.content.iter_any():
                try:
                    progress_data = json.loads(chunk)
                    if "progress" in progress_data:
                        task["progress"] = progress_data["progress"]
                except:
                    continue

async def call_wan_ti2v(wan_endpoint: str, task: Dict):
    """Call Wan TI2V endpoint"""
    url = f"{wan_endpoint}/api/ti2v"
    params = task["params"]
    
    # Add required parameters
    payload = {
        "prompt": params.get("prompt", ""),
        "image_url": params.get("image_url"),
        "duration": params.get("duration", 5),
        "fps": params.get("fps", 24),
        "style": params.get("style", "realistic"),
        "task_id": task["task_id"]
    }
    
    # Make the request
    async with aiohttp.ClientSession() as session:
        async with session.post(url, json=payload) as response:
            if response.status != 200:
                error = await response.text()
                raise Exception(f"Wan API error: {error}")
            
            # Stream the response to update progress
            async for chunk in response.content.iter_any():
                try:
                    progress_data = json.loads(chunk)
                    if "progress" in progress_data:
                        task["progress"] = progress_data["progress"]
                except:
                    continue

if __name__ == "__main__":
    import uvicorn
    import logging
    
    # Configure logging
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.StreamHandler(),
            logging.FileHandler("app.log")
        ]
    )
    
    # Start the server
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        reload=True,
        log_level="info"
    )
